import torch

from refiners.conversion.utils import Conversion, Hub, WeightRecipe

diffusers_recipe = WeightRecipe(
    key_map={
        "image_proj.proj": "image_proj.Linear",
        "image_proj.norm": "image_proj.LayerNorm",
        "ip_adapter.1.to_k_ip": "ip_adapter.000.to_k_ip",
        "ip_adapter.1.to_v_ip": "ip_adapter.000.to_v_ip",
        "ip_adapter.3.to_k_ip": "ip_adapter.001.to_k_ip",
        "ip_adapter.3.to_v_ip": "ip_adapter.001.to_v_ip",
        "ip_adapter.5.to_k_ip": "ip_adapter.002.to_k_ip",
        "ip_adapter.5.to_v_ip": "ip_adapter.002.to_v_ip",
        "ip_adapter.7.to_k_ip": "ip_adapter.003.to_k_ip",
        "ip_adapter.7.to_v_ip": "ip_adapter.003.to_v_ip",
        "ip_adapter.9.to_k_ip": "ip_adapter.004.to_k_ip",
        "ip_adapter.9.to_v_ip": "ip_adapter.004.to_v_ip",
        "ip_adapter.11.to_k_ip": "ip_adapter.005.to_k_ip",
        "ip_adapter.11.to_v_ip": "ip_adapter.005.to_v_ip",
        "ip_adapter.31.to_k_ip": "ip_adapter.006.to_k_ip",
        "ip_adapter.31.to_v_ip": "ip_adapter.006.to_v_ip",
        "ip_adapter.13.to_k_ip": "ip_adapter.007.to_k_ip",
        "ip_adapter.13.to_v_ip": "ip_adapter.007.to_v_ip",
        "ip_adapter.15.to_k_ip": "ip_adapter.008.to_k_ip",
        "ip_adapter.15.to_v_ip": "ip_adapter.008.to_v_ip",
        "ip_adapter.17.to_k_ip": "ip_adapter.009.to_k_ip",
        "ip_adapter.17.to_v_ip": "ip_adapter.009.to_v_ip",
        "ip_adapter.19.to_k_ip": "ip_adapter.010.to_k_ip",
        "ip_adapter.19.to_v_ip": "ip_adapter.010.to_v_ip",
        "ip_adapter.21.to_k_ip": "ip_adapter.011.to_k_ip",
        "ip_adapter.21.to_v_ip": "ip_adapter.011.to_v_ip",
        "ip_adapter.23.to_k_ip": "ip_adapter.012.to_k_ip",
        "ip_adapter.23.to_v_ip": "ip_adapter.012.to_v_ip",
        "ip_adapter.25.to_k_ip": "ip_adapter.013.to_k_ip",
        "ip_adapter.25.to_v_ip": "ip_adapter.013.to_v_ip",
        "ip_adapter.27.to_k_ip": "ip_adapter.014.to_k_ip",
        "ip_adapter.27.to_v_ip": "ip_adapter.014.to_v_ip",
        "ip_adapter.29.to_k_ip": "ip_adapter.015.to_k_ip",
        "ip_adapter.29.to_v_ip": "ip_adapter.015.to_v_ip",
    },
)

diffusers_plus_recipe = WeightRecipe(
    key_map={
        "image_proj.latents": "image_proj.LatentsToken.Parameter.weight",
        "image_proj.proj_in": "image_proj.Linear_1",
        "image_proj.proj_out": "image_proj.Linear_2",
        "image_proj.norm_out": "image_proj.LayerNorm",
        "image_proj.layers.0.0.norm1": "image_proj.Transformer.TransformerLayer_1.Residual_1.PerceiverAttention.Distribute.LayerNorm_1",
        "image_proj.layers.0.0.norm2": "image_proj.Transformer.TransformerLayer_1.Residual_1.PerceiverAttention.Distribute.LayerNorm_2",
        "image_proj.layers.0.0.to_q": "image_proj.Transformer.TransformerLayer_1.Residual_1.PerceiverAttention.Parallel.Chain_2.Linear",
        "image_proj.layers.0.0.to_kv": "image_proj.Transformer.TransformerLayer_1.Residual_1.PerceiverAttention.Parallel.Chain_1.Linear",
        "image_proj.layers.0.0.to_out": "image_proj.Transformer.TransformerLayer_1.Residual_1.PerceiverAttention.Linear",
        "image_proj.layers.0.1.0": "image_proj.Transformer.TransformerLayer_1.Residual_2.LayerNorm",
        "image_proj.layers.0.1.1": "image_proj.Transformer.TransformerLayer_1.Residual_2.FeedForward.Linear_1",
        "image_proj.layers.0.1.3": "image_proj.Transformer.TransformerLayer_1.Residual_2.FeedForward.Linear_2",
        "image_proj.layers.1.0.norm1": "image_proj.Transformer.TransformerLayer_2.Residual_1.PerceiverAttention.Distribute.LayerNorm_1",
        "image_proj.layers.1.0.norm2": "image_proj.Transformer.TransformerLayer_2.Residual_1.PerceiverAttention.Distribute.LayerNorm_2",
        "image_proj.layers.1.0.to_q": "image_proj.Transformer.TransformerLayer_2.Residual_1.PerceiverAttention.Parallel.Chain_2.Linear",
        "image_proj.layers.1.0.to_kv": "image_proj.Transformer.TransformerLayer_2.Residual_1.PerceiverAttention.Parallel.Chain_1.Linear",
        "image_proj.layers.1.0.to_out": "image_proj.Transformer.TransformerLayer_2.Residual_1.PerceiverAttention.Linear",
        "image_proj.layers.1.1.0": "image_proj.Transformer.TransformerLayer_2.Residual_2.LayerNorm",
        "image_proj.layers.1.1.1": "image_proj.Transformer.TransformerLayer_2.Residual_2.FeedForward.Linear_1",
        "image_proj.layers.1.1.3": "image_proj.Transformer.TransformerLayer_2.Residual_2.FeedForward.Linear_2",
        "image_proj.layers.2.0.norm1": "image_proj.Transformer.TransformerLayer_3.Residual_1.PerceiverAttention.Distribute.LayerNorm_1",
        "image_proj.layers.2.0.norm2": "image_proj.Transformer.TransformerLayer_3.Residual_1.PerceiverAttention.Distribute.LayerNorm_2",
        "image_proj.layers.2.0.to_q": "image_proj.Transformer.TransformerLayer_3.Residual_1.PerceiverAttention.Parallel.Chain_2.Linear",
        "image_proj.layers.2.0.to_kv": "image_proj.Transformer.TransformerLayer_3.Residual_1.PerceiverAttention.Parallel.Chain_1.Linear",
        "image_proj.layers.2.0.to_out": "image_proj.Transformer.TransformerLayer_3.Residual_1.PerceiverAttention.Linear",
        "image_proj.layers.2.1.0": "image_proj.Transformer.TransformerLayer_3.Residual_2.LayerNorm",
        "image_proj.layers.2.1.1": "image_proj.Transformer.TransformerLayer_3.Residual_2.FeedForward.Linear_1",
        "image_proj.layers.2.1.3": "image_proj.Transformer.TransformerLayer_3.Residual_2.FeedForward.Linear_2",
        "image_proj.layers.3.0.norm1": "image_proj.Transformer.TransformerLayer_4.Residual_1.PerceiverAttention.Distribute.LayerNorm_1",
        "image_proj.layers.3.0.norm2": "image_proj.Transformer.TransformerLayer_4.Residual_1.PerceiverAttention.Distribute.LayerNorm_2",
        "image_proj.layers.3.0.to_q": "image_proj.Transformer.TransformerLayer_4.Residual_1.PerceiverAttention.Parallel.Chain_2.Linear",
        "image_proj.layers.3.0.to_kv": "image_proj.Transformer.TransformerLayer_4.Residual_1.PerceiverAttention.Parallel.Chain_1.Linear",
        "image_proj.layers.3.0.to_out": "image_proj.Transformer.TransformerLayer_4.Residual_1.PerceiverAttention.Linear",
        "image_proj.layers.3.1.0": "image_proj.Transformer.TransformerLayer_4.Residual_2.LayerNorm",
        "image_proj.layers.3.1.1": "image_proj.Transformer.TransformerLayer_4.Residual_2.FeedForward.Linear_1",
        "image_proj.layers.3.1.3": "image_proj.Transformer.TransformerLayer_4.Residual_2.FeedForward.Linear_2",
        "ip_adapter.1.to_k_ip": "ip_adapter.000.to_k_ip",
        "ip_adapter.1.to_v_ip": "ip_adapter.000.to_v_ip",
        "ip_adapter.3.to_k_ip": "ip_adapter.001.to_k_ip",
        "ip_adapter.3.to_v_ip": "ip_adapter.001.to_v_ip",
        "ip_adapter.5.to_k_ip": "ip_adapter.002.to_k_ip",
        "ip_adapter.5.to_v_ip": "ip_adapter.002.to_v_ip",
        "ip_adapter.7.to_k_ip": "ip_adapter.003.to_k_ip",
        "ip_adapter.7.to_v_ip": "ip_adapter.003.to_v_ip",
        "ip_adapter.9.to_k_ip": "ip_adapter.004.to_k_ip",
        "ip_adapter.9.to_v_ip": "ip_adapter.004.to_v_ip",
        "ip_adapter.11.to_k_ip": "ip_adapter.005.to_k_ip",
        "ip_adapter.11.to_v_ip": "ip_adapter.005.to_v_ip",
        "ip_adapter.31.to_k_ip": "ip_adapter.006.to_k_ip",
        "ip_adapter.31.to_v_ip": "ip_adapter.006.to_v_ip",
        "ip_adapter.13.to_k_ip": "ip_adapter.007.to_k_ip",
        "ip_adapter.13.to_v_ip": "ip_adapter.007.to_v_ip",
        "ip_adapter.15.to_k_ip": "ip_adapter.008.to_k_ip",
        "ip_adapter.15.to_v_ip": "ip_adapter.008.to_v_ip",
        "ip_adapter.17.to_k_ip": "ip_adapter.009.to_k_ip",
        "ip_adapter.17.to_v_ip": "ip_adapter.009.to_v_ip",
        "ip_adapter.19.to_k_ip": "ip_adapter.010.to_k_ip",
        "ip_adapter.19.to_v_ip": "ip_adapter.010.to_v_ip",
        "ip_adapter.21.to_k_ip": "ip_adapter.011.to_k_ip",
        "ip_adapter.21.to_v_ip": "ip_adapter.011.to_v_ip",
        "ip_adapter.23.to_k_ip": "ip_adapter.012.to_k_ip",
        "ip_adapter.23.to_v_ip": "ip_adapter.012.to_v_ip",
        "ip_adapter.25.to_k_ip": "ip_adapter.013.to_k_ip",
        "ip_adapter.25.to_v_ip": "ip_adapter.013.to_v_ip",
        "ip_adapter.27.to_k_ip": "ip_adapter.014.to_k_ip",
        "ip_adapter.27.to_v_ip": "ip_adapter.014.to_v_ip",
        "ip_adapter.29.to_k_ip": "ip_adapter.015.to_k_ip",
        "ip_adapter.29.to_v_ip": "ip_adapter.015.to_v_ip",
    },
    tensor_reshapes={"image_proj.LatentsToken.Parameter.weight": (16, 768)},
)

base = Conversion(
    original=Hub(
        repo_id="h94/IP-Adapter",
        filename="models/ip-adapter_sd15.safetensors",
        revision="018e402774aeeddd60609b4ecdb7e298259dc729",
        expected_sha256="289b45f16d043d0bf542e45831f971dcdaabe18b656f11e86d9dfba7e9ee3369",
    ),
    converted=Hub(
        repo_id="refiners/sd15.ip_adapter",
        filename="model.safetensors",
        expected_sha256="ebabe531bac205e2fac942b353c22066abfb115b02f7fb72cd0e3361ee838ef3",
    ),
    recipe=diffusers_recipe,
    dtype=torch.float16,
)
plus = Conversion(
    original=Hub(
        repo_id="h94/IP-Adapter",
        filename="models/ip-adapter-plus_sd15.safetensors",
        revision="018e402774aeeddd60609b4ecdb7e298259dc729",
        expected_sha256="a1c250be40455cc61a43da1201ec3f1edaea71214865fb47f57927e06cbe4996",
    ),
    converted=Hub(
        repo_id="refiners/sd15.ip_adapter.plus",
        filename="model.safetensors",
        expected_sha256="6eae5d2098fa83e3b8bf416fb46dd77a921ad044650f13890e8d41d7c84a71d2",
    ),
    recipe=diffusers_plus_recipe,
    dtype=torch.float16,
)
